'''
Description: agibot_controller
Version: V1.0
Author: zw_1520@163.com
Date: 2024-11-19 15:51:41
LastEditors: zw_1520@163.com
LastEditTime: 2025-05-17 08:50:03
Copyright (C) All rights reserved.
'''

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSHistoryPolicy, QoSProfile, QoSReliabilityPolicy
from sensor_msgs.msg import JointState
from rclpy.executors import MultiThreadedExecutor
from std_msgs.msg import String
from std_msgs.msg import Int32
from robot_interface.srv import SetRobotTask
from robot_interface.msg import RobotTaskInfo
import os
from ament_index_python.packages import get_package_prefix
from agibot_controller.play_motion import *
from agibot_controller.arm import *
from agibot_controller.navigation import *
from agibot_controller.play_tts import *
from agibot_controller.play_motion_by_http import *
from agibot_controller.navigation_by_http import *
from agibot_controller.arm_joint_control import stop_agi_motion_playing


class AgibotController(Node):
    def __init__(self):
        super().__init__("agibot_controller")
        package_name = "agibot_controller"
        self.package_install_directory = get_package_prefix(package_name)
        self.get_logger().info(
            f"Package install directory: {self.package_install_directory}")

        self.subscription = self.create_subscription(
            RobotTaskInfo, "/agibot_task_control", self.topic_callback, 10
        )
        self.subscription

        qos_profile = QoSProfile(
            history=QoSHistoryPolicy.KEEP_LAST, depth=10, reliability=QoSReliabilityPolicy.BEST_EFFORT
        )

        self.arm_joint_subscription = self.create_subscription(
            JointState, '/motion/control/arm_joint_state', self.joint_listener_callback, qos_profile)
        self.arm_joint_subscription

        self.arm_joint_publisher_ = self.create_publisher(
            JointState, '/motion/control/arm_joint_command', qos_profile)
        self.arm_joint_publisher_

        self.get_logger().info("agibot controller initialized.")

        self.agi_control_service = self.create_service(
            SetRobotTask, '/agibot_control_service', self.agi_control_service_callback)
        self.get_logger().info("Agibot control service start.")
        self.agi_control_service

        # Initialize zero position
        self.current_joint_position = []
        self.pos_ini = [0.0, 1.3, 0.0, -0.1, 1.57, 0.0,
                        0.0, 0.0, -1.3, 0.0, 0.1, 1.57, 0.0, 0.0]
        self.move_count = 200
        self.step = 1.4 / self.move_count * (-1)
        self.bow_step = 1.57 / self.move_count * (-1)
        self.is_joint_state_subscribed = False
        self.is_leftarm_lifted = False
        self.get_logger().info("agibot controller initialized.")

        # Starts thread to run new thread fun
        # self.thread = threading.Thread(target=self.new_thread_fun)
        # self.thread.start()
        # self.get_logger().info("new thread started.")

    def new_thread_fun(self):
        # Waits for the joint state
        target_joint_state = self.pos_ini
        target_joint_state[1] = 0.0
        target_joint_state[4] = 0.0
        while not self.is_joint_state_subscribed:
            self.get_logger().info("Waiting for joint state")
            time.sleep(1)
        self.move_arm_to_target_pose(target_joint_state)

    def recover_left_arm_thread_fun(self):
        # Waits for the joint state
        target_joint_state = self.pos_ini
        while not self.is_joint_state_subscribed:
            self.get_logger().info("Waiting for joint state")
            time.sleep(1)
        self.move_arm_to_target_pose(target_joint_state)

    def joint_listener_callback(self, msg: JointState):
        if not self.is_joint_state_subscribed:
            self.get_logger().info("joint state subscribed.")
            self.get_logger().info(f"robot arm position: {msg.position}")
        self.is_joint_state_subscribed = True
        self.current_joint_position = msg.position

    def topic_callback(self, msg):
        if (
            msg.task_type == "play_motion"
        ):
            try:
                self.get_logger().info(f"Receive command: {msg.task_type}")
                file_path = os.path.join(
                    self.package_install_directory, "cfg", "play_motion.yaml")
                self.get_logger().info(
                    f"Play motion config file path: {file_path}")
                self.get_logger().info(
                    f"Play motion file: {msg.task_param}")
                (msg.task_param, 10000)
                if msg.task_param == "welcome":
                    play_motion_by_http(msg.task_param, 10000)

                elif msg.task_param == "lift_left_arm":
                    if self.is_leftarm_lifted:
                        self.get_logger().info()("Left arm already lifted")
                        return
                    stop_agi_motion_playing()
                    self.thread = threading.Thread(target=self.new_thread_fun)
                    self.thread.start()
                    self.get_logger().info("new thread started.")
                    self.is_leftarm_lifted = True
                    self.get_logger().info("Left arm lifted")

                elif msg.task_param == "put_down_left_arm":
                    if change_robot_to_navigation_mode():
                        self.is_leftarm_lifted = False
                        self.get_logger().info("Left arm put down successfully.")
                    else:
                        self.get_logger().error(f"Failed to put down left arm.")

            except Exception as e:
                self.get_logger().error(f"Failed to control: {e}")
        elif msg.task_type == "move_arm":
            try:
                self.get_logger().info(f"Receive command: {msg.task_type}")
                file_path = os.path.join(
                    self.package_install_directory, "cfg", "arm.yaml")
                self.get_logger().info(
                    f"Arm control config file path: {file_path}")
                arm_control(file_path)
            except Exception as e:
                self.get_logger().error(f"Failed to control arm: {e}")
        elif msg.task_type == "navigation":
            try:
                self.get_logger().info(f"Receive command: {msg.task_type}")
                file_path = os.path.join(
                    self.package_install_directory, "cfg", "navigation.yaml")
                self.get_logger().info(
                    f"Navigation config file path: {file_path}")
                # navi_to_target(file_path)
                print(msg.task_param)

                # Parsers point id
                data = json.loads(msg.task_param)

                if "point_id" not in data:
                    self.get_logger().error("Invalid point id")
                    return
                point_id = data["point_id"]

                navigation_by_http(point_id)
            except Exception as e:
                self.get_logger().error(f"Failed to navi: {e}")
        elif msg.task_type == "tts":
            try:
                self.get_logger().info(f"Receive command: {msg.task_type}")
                init_tts()
                play_tts(msg.task_param)
            except Exception as e:
                self.get_logger().error(f"Failed to navi: {e}")

    def agi_control_service_callback(self, request, response):
        if (
            request.task_type == "play_motion"
        ):
            try:
                self.get_logger().info(f"Receive command: {request.task_type}")
                file_path = os.path.join(
                    self.package_install_directory, "cfg", "play_motion.yaml")
                self.get_logger().info(
                    f"Play motion config file path: {file_path}")
                # play_motion_by_external(
                #     file_path, "/agibot/data/var/rc/motion_player/default/投篮/投篮", 10000)
                response.success = True
            except Exception as e:
                response.success = False
                self.get_logger().error(f"Failedxxx to control: {e}")
        elif request.task_type == "move_arm":
            try:
                self.get_logger().info(f"Receive command: {request.task_type}")
                file_path = os.path.join(
                    self.package_install_directory, "cfg", "arm.yaml")
                self.get_logger().info(
                    f"Arm control config file path: {file_path}")
                arm_control(file_path)
                response.success = True
            except Exception as e:
                self.get_logger().error(f"Failed to control arm: {e}")
        elif request.task_type == "navigation":
            try:
                self.get_logger().info(f"Receive command: {request.task_type}")
                file_path = os.path.join(
                    self.package_install_directory, "cfg", "navigation.yaml")
                self.get_logger().info(
                    f"Navigation config file path: {file_path}")
                response.success = True
            except Exception as e:
                self.get_logger().error(f"Failed to navi: {e}")

    def move_arm_to_target_pose(self, target_pose):
        self.get_logger().info(f"Move arm to target pose: {target_pose}")
        current_joint_position = self.current_joint_position
        delter_joint_position = []
        for i in range(len(self.current_joint_position)):
            delter_joint_position.append(
                target_pose[i] - current_joint_position[i])
        self.get_logger().info(
            f"delter_joint_position: {delter_joint_position}")

        # Calculate the joint move step
        joint_move_step = []
        for i in range(len(current_joint_position)):
            joint_move_step.append(
                delter_joint_position[i] / self.move_count)
        self.get_logger().info(
            f"joint_move_step: {joint_move_step}")

        step_control_pose = current_joint_position
        for i in range(self.move_count):
            for j in range(len(step_control_pose)):
                step_control_pose[j] = step_control_pose[j] + \
                    joint_move_step[j]
            self.get_logger().info(
                f"step_control_pose {i}: {step_control_pose}")
            # Publishes the joint positions to the joint_command topic
            msg = JointState()
            msg.header.stamp = self.get_clock().now().to_msg()
            msg.position = step_control_pose
            self.arm_joint_publisher_.publish(msg)
            time.sleep(0.01)


def main(args=None):
    # rclpy.init(args=args)
    # sbr = AgibotController()
    # rclpy.spin(sbr)
    # rclpy.shutdown()

    rclpy.init(args=args)
    agibot_controller_node = AgibotController()
    executor = MultiThreadedExecutor()
    executor.add_node(agibot_controller_node)
    executor.spin()
    agibot_controller_node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
